#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将FrankMocap输出的所有pkl文件合并成一个pkl文件
支持两种输出格式：
1. frame_list: 保存为帧列表，每个元素是原始的FrankMocap输出格式
2. animation_data: 保存为类似npz的结构但用pkl格式，包含完整的动画轨迹数据

轨迹数据包括：
- 顶点轨迹 (vertices) - 3D顶点位置
- 图像空间关节轨迹 (body_joints_img, hand_joints_img)
- 手部朝向轨迹 (hand_orient)
- 计算的轨迹数据 (root_center_trajectory, velocity, 等)
- 相机和边界框轨迹
"""

import os
import pickle
import numpy as np
from tqdm import tqdm
import argparse

try:
    from scipy.interpolate import interp1d
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False
    print("Warning: scipy not available, will use simple interpolation")

def load_pkl_frame_original(pkl_path):
    """加载单个pkl文件，保持原始结构"""
    try:
        with open(pkl_path, 'rb') as f:
            data = pickle.load(f)
        return data
    except Exception as e:
        print(f"Error loading {pkl_path}: {e}")
        return None

def load_pkl_frame_processed(pkl_path):
    """加载单个pkl文件并提取关键数据（用于animation_data格式）"""
    try:
        with open(pkl_path, 'rb') as f:
            data = pickle.load(f)

        if 'pred_output_list' not in data or len(data['pred_output_list']) == 0:
            return None

        pred_output = data['pred_output_list'][0]

        # 提取所有需要的数据
        frame_data = {}

        # 基础信息
        frame_data['demo_type'] = data.get('demo_type', '')
        frame_data['smpl_type'] = data.get('smpl_type', '')
        
        # 用于标记哪些数据有效
        frame_data['valid_data'] = {}

        # 定义需要插值的数据字段 - 基于实际存在的字段
        interpolatable_fields = [
            # 基础姿态数据（SMPL参数）
            ('pred_body_pose', 'body_pose'),              # (1, 72) - 身体姿态参数
            ('pred_betas', 'betas'),                      # (1, 10) - 形状参数
            
            # 手部姿态数据  
            ('pred_left_hand_pose', 'left_hand_pose'),    # (1, 45) - 左手姿态
            ('pred_right_hand_pose', 'right_hand_pose'),  # (1, 45) - 右手姿态
            ('pred_left_hand_betas', 'left_hand_betas'),  # (1, 10) - 左手形状
            ('pred_right_hand_betas', 'right_hand_betas'), # (1, 10) - 右手形状
            
            # 几何数据（实际轨迹数据）
            ('pred_vertices_smpl', 'vertices'),           # (10475, 3) - 3D顶点轨迹
            ('pred_vertices_img', 'vertices_img'),        # (10475, 3) - 图像空间顶点
            
            # 相机数据
            ('pred_camera', 'camera'),                    # (3,) - 相机参数
            
            # 图像空间关节轨迹（主要轨迹数据）
            ('pred_body_joints_img', 'body_joints_img'),    # (49, 3) - 身体关节图像轨迹
            ('pred_lhand_joints_img', 'left_hand_joints_img'), # (21, 3) - 左手关节图像轨迹
            ('pred_rhand_joints_img', 'right_hand_joints_img'), # (21, 3) - 右手关节图像轨迹
        ]

        # 处理可插值的数据字段
        for source_key, target_key in interpolatable_fields:
            if source_key in pred_output and pred_output[source_key] is not None:
                frame_data[target_key] = pred_output[source_key].astype(np.float32)
                frame_data['valid_data'][target_key] = True
            else:
                frame_data['valid_data'][target_key] = False

        # 面信息 (只需要一次，所有帧都相同)
        if 'faces' in pred_output and pred_output['faces'] is not None:
            frame_data['faces'] = pred_output['faces']

        # 手部方向信息 (也需要插值)
        hand_orient_fields = [
            'left_hand_local_orient_body', 'left_hand_global_orient_body',
            'right_hand_local_orient_body', 'right_hand_global_orient_body',
            'left_hand_local_orient_hand', 'left_hand_global_orient_hand',
            'right_hand_local_orient_hand', 'right_hand_global_orient_hand'
        ]
        
        for key in hand_orient_fields:
            if key in pred_output and pred_output[key] is not None:
                frame_data[key] = pred_output[key].astype(np.float32)
                frame_data['valid_data'][key] = True
            else:
                frame_data['valid_data'][key] = False

        # 边界框和缩放信息（轨迹相关）
        bbox_scale_fields = [
            ('bbox_scale_ratio', 'bbox_scale_ratio'),
            ('bbox_top_left', 'bbox_top_left'),
            ('left_hand_bbox_scale_ratio', 'left_hand_bbox_scale_ratio'),
            ('left_hand_bbox_top_left', 'left_hand_bbox_top_left'),
            ('right_hand_bbox_scale_ratio', 'right_hand_bbox_scale_ratio'),
            ('right_hand_bbox_top_left', 'right_hand_bbox_top_left'),
            ('left_hand_camera', 'left_hand_camera'),
            ('right_hand_camera', 'right_hand_camera'),
        ]
        
        for source_key, target_key in bbox_scale_fields:
            if source_key in pred_output and pred_output[source_key] is not None:
                if isinstance(pred_output[source_key], np.ndarray):
                    frame_data[target_key] = pred_output[source_key].astype(np.float32)
                    frame_data['valid_data'][target_key] = True
                elif isinstance(pred_output[source_key], (int, float)):
                    frame_data[target_key] = np.array([pred_output[source_key]], dtype=np.float32)
                    frame_data['valid_data'][target_key] = True
                else:
                    frame_data['valid_data'][target_key] = False
            else:
                frame_data['valid_data'][target_key] = False

        return frame_data

    except Exception as e:
        error_msg = str(e)
        if "'NoneType' object has no attribute 'astype'" in error_msg:
            print(f"Error loading {pkl_path}: Data contains None values - this frame likely has incomplete motion capture data")
        else:
            print(f"Error loading {pkl_path}: {e}")
        return None

def interpolate_missing_data(animation_data, all_valid_flags):
    """
    对缺失的数据进行插值处理
    
    Args:
        animation_data: 包含所有动画数据的字典
        all_valid_flags: 每帧每个字段的有效性标记
    """
    num_frames = len(all_valid_flags)
    
    print("Interpolating missing data...")
    interpolated_count = 0
    
    # 获取所有需要插值的字段
    all_fields = set()
    for frame_flags in all_valid_flags:
        if frame_flags:
            all_fields.update(frame_flags.keys())
    
    for field in all_fields:
        if field not in animation_data:
            continue
            
        # 收集有效帧的索引和数据
        valid_indices = []
        valid_data = []
        
        for frame_idx, frame_flags in enumerate(all_valid_flags):
            if frame_flags and frame_flags.get(field, False):
                valid_indices.append(frame_idx)
                valid_data.append(animation_data[field][frame_idx])
        
        if len(valid_indices) < 2:
            # 如果有效数据太少，使用第一个有效值填充所有缺失帧
            if len(valid_indices) == 1:
                fill_value = valid_data[0]
                for frame_idx in range(num_frames):
                    frame_flags = all_valid_flags[frame_idx]
                    if not (frame_flags and frame_flags.get(field, False)):
                        animation_data[field][frame_idx] = fill_value
                        interpolated_count += 1
            continue
            
        # 对多维数据进行插值
        valid_indices = np.array(valid_indices)
        valid_data = np.array(valid_data)
        
        # 找到需要插值的帧
        for frame_idx in range(num_frames):
            frame_flags = all_valid_flags[frame_idx]
            if not (frame_flags and frame_flags.get(field, False)):
                # 进行插值
                if frame_idx < valid_indices[0]:
                    # 前向填充
                    animation_data[field][frame_idx] = valid_data[0]
                elif frame_idx > valid_indices[-1]:
                    # 后向填充
                    animation_data[field][frame_idx] = valid_data[-1]
                else:
                    # 线性插值
                    if HAS_SCIPY:
                        # 使用scipy进行精确插值
                        interp_func = interp1d(valid_indices, valid_data, axis=0, 
                                              kind='linear', bounds_error=False, 
                                              fill_value=(valid_data[0], valid_data[-1]))
                        animation_data[field][frame_idx] = interp_func(frame_idx)
                    else:
                        # 简单线性插值
                        # 找到前后最近的两个有效帧
                        before_idx = np.where(valid_indices < frame_idx)[0]
                        after_idx = np.where(valid_indices > frame_idx)[0]
                        
                        if len(before_idx) > 0 and len(after_idx) > 0:
                            before = before_idx[-1]  # 最近的前一个
                            after = after_idx[0]     # 最近的后一个
                            
                            t = (frame_idx - valid_indices[before]) / (valid_indices[after] - valid_indices[before])
                            interpolated = valid_data[before] * (1 - t) + valid_data[after] * t
                            animation_data[field][frame_idx] = interpolated
                        else:
                            # 如果只有一边有数据，使用最近的有效值
                            if len(before_idx) > 0:
                                animation_data[field][frame_idx] = valid_data[before_idx[-1]]
                            else:
                                animation_data[field][frame_idx] = valid_data[after_idx[0]]
                interpolated_count += 1
    
    if interpolated_count > 0:
        print(f"Interpolated {interpolated_count} missing data points")
    
    return animation_data

def compute_trajectory_data(animation_data):
    """
    从基础动画数据计算额外的轨迹信息
    基于实际存在的字段：图像空间关节、顶点等
    """
    print("Computing additional trajectory data from available fields...")
    
    num_frames = animation_data.get('num_frames', 0)
    if num_frames == 0:
        return animation_data
    
    # 1. 从顶点数据计算根部中心轨迹
    if 'vertices' in animation_data:
        vertices = animation_data['vertices']  # (frames, 10475, 3)
        if vertices.shape[0] > 0:
            # 计算顶点中心作为根部位置轨迹
            root_center = np.mean(vertices, axis=1)  # (frames, 3)
            animation_data['root_center_trajectory'] = root_center
            
            # 计算根部运动轨迹的速度和加速度
            velocity = np.zeros_like(root_center)
            velocity[1:] = root_center[1:] - root_center[:-1]
            animation_data['root_velocity_trajectory'] = velocity
            
            acceleration = np.zeros_like(velocity)
            acceleration[1:] = velocity[1:] - velocity[:-1]
            animation_data['root_acceleration_trajectory'] = acceleration
            
            # 累积路径长度
            path_length = np.zeros(num_frames)
            for i in range(1, num_frames):
                path_length[i] = path_length[i-1] + np.linalg.norm(root_center[i] - root_center[i-1])
            animation_data['cumulative_path_length'] = path_length
    
    # 2. 从图像空间身体关节计算轨迹
    if 'body_joints_img' in animation_data:
        body_joints = animation_data['body_joints_img']  # (frames, 49, 3)
        if body_joints.shape[0] > 0:
            # 身体中心轨迹（在图像空间）
            body_center = np.mean(body_joints, axis=1)  # (frames, 3)
            animation_data['body_center_trajectory'] = body_center
            
            # 关键关节轨迹（如果有足够的关节）
            if body_joints.shape[1] >= 25:
                # SMPL关键关节索引（根据SMPL定义）
                key_joints_indices = [0, 1, 2, 4, 5, 16, 17, 18, 19]  # 主要关节
                valid_indices = [i for i in key_joints_indices if i < body_joints.shape[1]]
                if valid_indices:
                    animation_data['key_joints_trajectory'] = body_joints[:, valid_indices, :]
            
            # 身体关节运动速度
            body_velocity = np.zeros_like(body_joints)
            body_velocity[1:] = body_joints[1:] - body_joints[:-1]
            animation_data['body_joints_velocity'] = body_velocity
    
    # 3. 手部轨迹中心（图像空间）
    for hand_name, hand_key in [('left', 'left_hand_joints_img'), ('right', 'right_hand_joints_img')]:
        if hand_key in animation_data:
            hand_joints = animation_data[hand_key]  # (frames, 21, 3)
            if hand_joints.shape[0] > 0:
                # 手部中心轨迹
                hand_center = np.mean(hand_joints, axis=1)  # (frames, 3)
                animation_data[f'{hand_name}_hand_center_trajectory'] = hand_center
                
                # 手部运动速度
                hand_velocity = np.zeros_like(hand_center)
                hand_velocity[1:] = hand_center[1:] - hand_center[:-1]
                animation_data[f'{hand_name}_hand_velocity'] = hand_velocity
                
                # 手指关节轨迹（关键点）
                if hand_joints.shape[1] >= 21:  # 标准手部关节数
                    # 取几个关键手指关节
                    finger_tip_indices = [4, 8, 12, 16, 20]  # 拇指到小指的指尖
                    animation_data[f'{hand_name}_fingertips_trajectory'] = hand_joints[:, finger_tip_indices, :]
    
    # 4. 从手部朝向信息计算朝向轨迹
    hand_orient_keys = [
        'left_hand_global_orient_body', 'right_hand_global_orient_body',
        'left_hand_local_orient_body', 'right_hand_local_orient_body'
    ]
    
    for orient_key in hand_orient_keys:
        if orient_key in animation_data:
            orient_data = animation_data[orient_key]  # (frames, 1, 3)
            if orient_data.shape[0] > 0:
                # 朝向变化轨迹
                orient_velocity = np.zeros_like(orient_data)
                orient_velocity[1:] = orient_data[1:] - orient_data[:-1]
                animation_data[f'{orient_key}_velocity'] = orient_velocity
    
    # 5. 相机参数轨迹（如果相机在移动）
    if 'camera' in animation_data:
        camera = animation_data['camera']  # (frames, 3)
        if camera.shape[0] > 0:
            # 相机运动轨迹
            camera_velocity = np.zeros_like(camera)
            camera_velocity[1:] = camera[1:] - camera[:-1]
            animation_data['camera_motion_trajectory'] = camera_velocity
    
    # 6. 身体姿态参数变化轨迹
    if 'body_pose' in animation_data:
        body_pose = animation_data['body_pose']  # (frames, 1, 72)
        if body_pose.shape[0] > 0:
            # 姿态变化速度
            pose_velocity = np.zeros_like(body_pose)
            pose_velocity[1:] = body_pose[1:] - body_pose[:-1]
            animation_data['pose_change_trajectory'] = pose_velocity
    
    print(f"Added trajectory data based on available fields:")
    print(f"  - Root center trajectory from vertices")
    print(f"  - Body joints trajectory (image space)")
    print(f"  - Hand center trajectories")
    print(f"  - Velocity and acceleration trajectories")
    print(f"  - Camera motion trajectory")
    
    return animation_data

def merge_pkls_to_pkl_frame_list(mocap_dir, output_path, frame_range=None):
    """
    将所有pkl文件合并成一个pkl文件 - 帧列表格式
    输出格式：{'frames': [frame1_data, frame2_data, ...], 'metadata': {...}}
    """
    # 获取所有pkl文件
    pkl_files = sorted([f for f in os.listdir(mocap_dir) if f.endswith('.pkl')])

    if not pkl_files:
        print(f"No pkl files found in {mocap_dir}")
        return False

    print(f"Found {len(pkl_files)} pkl files")

    # 应用帧范围过滤
    if frame_range:
        start_frame, end_frame = frame_range
        if end_frame is not None:
            pkl_files = pkl_files[start_frame:end_frame]
        else:
            pkl_files = pkl_files[start_frame:]
        print(f"Processing frames {start_frame} to {end_frame if end_frame is not None else len(pkl_files)+start_frame-1} ({len(pkl_files)} files)")

    # 加载所有帧
    print("Loading all frames...")
    frames = []
    skipped_count = 0
    
    for i, pkl_file in enumerate(tqdm(pkl_files, desc="Loading frames")):
        pkl_path = os.path.join(mocap_dir, pkl_file)
        frame_data = load_pkl_frame_original(pkl_path)
        
        if frame_data is None:
            print(f"Skipping frame {i}: {pkl_file}")
            skipped_count += 1
            continue
            
        # 添加帧索引信息
        frame_data['frame_index'] = i
        frame_data['original_filename'] = pkl_file
        frames.append(frame_data)

    # 准备输出数据
    output_data = {
        'frames': frames,
        'metadata': {
            'num_frames': len(frames),
            'skipped_frames': skipped_count,
            'fps': 30.0,
            'total_duration': len(frames) / 30.0,
            'format': 'frame_list',
            'original_file_count': len(pkl_files) + skipped_count
        }
    }

    # 添加第一帧的基础信息到元数据
    if frames:
        first_frame = frames[0]
        output_data['metadata']['demo_type'] = first_frame.get('demo_type', '')
        output_data['metadata']['smpl_type'] = first_frame.get('smpl_type', '')

    # 创建输出目录
    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    # 保存为pkl文件
    print(f"Saving animation data to {output_path}")
    try:
        with open(output_path, 'wb') as f:
            pickle.dump(output_data, f, protocol=2)  # 兼容老版本Blender
    except Exception as e:
        print(f"Error saving pkl file: {e}")
        return False

    # 打印摘要信息
    print("\n=== Animation Data Summary ===")
    print(f"Total frames: {len(frames)}")
    print(f"Skipped frames: {skipped_count}")
    print(f"Duration: {len(frames)/30.0:.2f} seconds")
    print(f"Format: Frame List")

    # 获取文件大小
    try:
        file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
        print(f"Output file size: {file_size:.2f} MB")
    except OSError:
        print("Could not determine file size")

    return True

def merge_pkls_to_pkl_animation_data(mocap_dir, output_path, frame_range=None):
    """
    将所有pkl文件合并成一个pkl文件 - 动画数据格式
    输出格式类似npz但使用pkl保存，支持插值
    """
    # 统计跳过的帧
    skipped_stats = {
        'file_error': 0,
        'missing_data': 0,
        'none_values': 0
    }

    # 获取所有pkl文件
    pkl_files = sorted([f for f in os.listdir(mocap_dir) if f.endswith('.pkl')])

    if not pkl_files:
        print(f"No pkl files found in {mocap_dir}")
        return False

    print(f"Found {len(pkl_files)} pkl files")

    # 应用帧范围过滤
    if frame_range:
        start_frame, end_frame = frame_range
        if end_frame is not None:
            pkl_files = pkl_files[start_frame:end_frame]
        else:
            pkl_files = pkl_files[start_frame:]
        print(f"Processing frames {start_frame} to {end_frame if end_frame is not None else len(pkl_files)+start_frame-1} ({len(pkl_files)} files)")

    # 初始化数据容器
    animation_data = {}
    faces = None  # 面信息对所有帧都相同
    all_valid_flags = []  # 保存每帧的有效数据标记
    
    # 先处理所有帧，收集数据结构信息
    print("Loading all frames...")
    all_frame_data = []
    for i, pkl_file in enumerate(tqdm(pkl_files, desc="Loading frames")):
        pkl_path = os.path.join(mocap_dir, pkl_file)
        frame_data = load_pkl_frame_processed(pkl_path)

        if frame_data is None:
            print(f"Skipping frame {i}: {pkl_file}")
            skipped_stats['file_error'] += 1
            # 创建空的frame_data以保持帧数
            frame_data = {'valid_data': {}}
            
        all_frame_data.append(frame_data)
        all_valid_flags.append(frame_data.get('valid_data', {}))

    num_frames = len(all_frame_data)
    print(f"Loaded {num_frames} frames, initializing data arrays...")

    # 首先确定所有可能的数据字段（从有效的帧中获取）
    all_data_fields = set()
    demo_type = smpl_type = ""
    
    for frame_data in all_frame_data:
        if 'demo_type' in frame_data:
            demo_type = frame_data['demo_type']
        if 'smpl_type' in frame_data:
            smpl_type = frame_data['smpl_type']
        if 'faces' in frame_data:
            faces = frame_data['faces']
            
        # 收集所有数据字段
        for key, value in frame_data.items():
            if key not in ['demo_type', 'smpl_type', 'faces', 'valid_data'] and isinstance(value, np.ndarray):
                all_data_fields.add(key)
    
    # 设置基础信息
    animation_data['demo_type'] = demo_type
    animation_data['smpl_type'] = smpl_type
    
    # 初始化所有数据数组
    print("Initializing data arrays...")
    for field in all_data_fields:
        # 找到第一个包含此字段的帧来确定数据形状
        for frame_data in all_frame_data:
            if field in frame_data:
                value = frame_data[field]
                shape = (num_frames,) + value.shape
                animation_data[field] = np.zeros(shape, dtype=value.dtype)
                break
    
    # 填充数据
    print("Filling data arrays...")
    for i, frame_data in enumerate(tqdm(all_frame_data, desc="Filling arrays")):
        for field in all_data_fields:
            if field in frame_data and field in animation_data:
                animation_data[field][i] = frame_data[field]

    # 添加面信息
    if faces is not None:
        animation_data['faces'] = faces

    # 进行插值处理以填充缺失数据
    animation_data = interpolate_missing_data(animation_data, all_valid_flags)
    
    # 计算额外的轨迹数据
    animation_data = compute_trajectory_data(animation_data)
    
    # 检查数据质量
    total_missing = sum(skipped_stats.values())
    if total_missing > 0:
        print(f"Warning: {total_missing} frames had data issues, but were interpolated")
    
    # 检查是否帧数太少
    if num_frames < 10:
        print(f"Warning: Very few frames ({num_frames}) - animation may be too short")

    # 添加元数据
    animation_data['num_frames'] = num_frames
    animation_data['fps'] = 30.0  # FrankMocap默认30fps
    animation_data['total_duration'] = num_frames / 30.0
    animation_data['format'] = 'animation_data'

    # 创建输出目录
    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    # 保存为pkl文件
    print(f"Saving animation data to {output_path}")
    try:
        with open(output_path, 'wb') as f:
            pickle.dump(animation_data, f, protocol=2)  # 兼容老版本Blender
    except Exception as e:
        print(f"Error saving pkl file: {e}")
        return False

    # 打印摘要信息
    print("\n=== Animation Data Summary ===")
    print(f"Total frames: {num_frames}")
    print(f"Duration: {num_frames/30.0:.2f} seconds")
    print(f"Format: Animation Data")
    
    # 显示数据质量统计
    if skipped_stats['file_error'] > 0:
        print(f"Frames with missing data: {skipped_stats['file_error']} ({skipped_stats['file_error']*100/num_frames:.1f}%)")
        print("Missing data was interpolated from neighboring frames")
    
    print(f"Data arrays saved:")

    for key, value in animation_data.items():
        if isinstance(value, np.ndarray):
            print(f"  {key}: {value.shape} ({value.dtype})")
        else:
            print(f"  {key}: {type(value).__name__} = {value}")

    # 获取文件大小
    try:
        file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
        print(f"Output file size: {file_size:.2f} MB")
    except OSError:
        print("Could not determine file size")

    return True

def main():
    parser = argparse.ArgumentParser(description='Merge FrankMocap pkl files into a single pkl file')
    parser.add_argument('--input', '-i',
                       default='./mocap_output_full/mocap',
                       help='Input directory containing pkl files')
    parser.add_argument('--output', '-o',
                       default='./output_pkl/merged_animation.pkl',
                       help='Output pkl file path')
    parser.add_argument('--format', '-f',
                       choices=['frame_list', 'animation_data'],
                       default='frame_list',
                       help='Output format: frame_list (original structure) or animation_data (processed with interpolation)')
    parser.add_argument('--start-frame', type=int, default=None,
                       help='Start frame index (0-based)')
    parser.add_argument('--end-frame', type=int, default=None,
                       help='End frame index (exclusive)')

    args = parser.parse_args()

    # 检查输入目录
    if not os.path.exists(args.input):
        print(f"Error: Input directory {args.input} does not exist")
        return

    # 设置帧范围
    frame_range = None
    if args.start_frame is not None or args.end_frame is not None:
        start = args.start_frame if args.start_frame is not None else 0
        end = args.end_frame  # None表示到最后一帧
        frame_range = (start, end)
        print(f"Frame range specified: {start} to {end if end is not None else 'end'}")

    # 执行合并
    print("=== FrankMocap PKL Merger ===")
    print(f"Input directory: {args.input}")
    print(f"Output file: {args.output}")
    print(f"Output format: {args.format}")

    if args.format == 'frame_list':
        success = merge_pkls_to_pkl_frame_list(args.input, args.output, frame_range)
        format_description = "Frame List (preserves original structure)"
    else:
        success = merge_pkls_to_pkl_animation_data(args.input, args.output, frame_range)
        format_description = "Animation Data (processed with interpolation)"

    if success:
        print("\n✅ Merge completed successfully!")
        print(f"Animation data saved to: {args.output}")
        print(f"Format: {format_description}")
        
        if args.format == 'frame_list':
            print("\nFormat details:")
            print("- Preserves original FrankMocap structure")
            print("- Each frame contains complete original data")
            print("- Easy to access individual frame data")
            print("- Contains all original trajectory fields")
            print("- Use: data = pickle.load(file); frames = data['frames']")
        else:
            print("\nFormat details:")
            print("- Optimized for animation processing")
            print("- Missing data interpolated for smooth animation")
            print("- Contains comprehensive trajectory data:")
            print("  * Root center trajectory (from vertex data)")
            print("  * Body joints trajectory (image space)")
            print("  * Hand center trajectories")
            print("  * Velocity and acceleration trajectories")
            print("  * Vertex trajectories (detailed 3D mesh)")
            print("  * Hand orientation trajectories")
            print("- Use: data = pickle.load(file); root_traj = data['root_center_trajectory']")
    else:
        print("\n❌ Merge failed!")

if __name__ == "__main__":
    main()